import numpy as np
from . import LocalSearchOperator, register_operator


class AdjacentSwapOperator(LocalSearchOperator):
    """相邻交换算子，只交换相邻的城市"""

    def __init__(self):
        self.name = "adjacent_swap"
        self.supported_problems = ["tsp"]

    def generate_neighborhood(self, solution, problem_instance):
        """生成相邻交换邻域"""
        neighborhood = []
        n = len(solution)

        # 生成所有相邻交换
        for i in range(n):
            j = (i + 1) % n  # 只与下一个城市交换
            new_solution = solution.copy()
            new_solution[i], new_solution[j] = new_solution[j], new_solution[i]
            neighborhood.append(new_solution)

        return neighborhood

    def efficient_adjacent_swap(self, solution, problem_instance, intensity=20):
        """高效的相邻交换实现"""
        n = len(solution)
        current_solution = solution.copy()
        current_cost = problem_instance.evaluate(current_solution)
        improved = True
        iterations = 0

        while improved and iterations < intensity:
            improved = False

            # 尝试所有相邻交换
            for i in range(n):
                j = (i + 1) % n

                # 计算交换前的相关边成本
                prev_i = (i - 1) % n
                next_j = (j + 1) % n

                # 计算交换前的成本
                old_cost = (
                        problem_instance.distance_matrix[current_solution[prev_i]][current_solution[i]] +
                        problem_instance.distance_matrix[current_solution[i]][current_solution[j]] +
                        problem_instance.distance_matrix[current_solution[j]][current_solution[next_j]]
                )

                # 执行交换
                new_solution = current_solution.copy()
                new_solution[i], new_solution[j] = new_solution[j], new_solution[i]

                # 计算交换后的成本
                new_cost = (
                        problem_instance.distance_matrix[new_solution[prev_i]][new_solution[i]] +
                        problem_instance.distance_matrix[new_solution[i]][new_solution[j]] +
                        problem_instance.distance_matrix[new_solution[j]][new_solution[next_j]]
                )

                # 计算成本变化
                cost_change = new_cost - old_cost

                # 如果成本降低，则接受交换
                if cost_change < 0:
                    current_solution = new_solution
                    current_cost += cost_change
                    improved = True
                    break  # 找到一个改进就跳出内层循环

            iterations += 1

        improvement = problem_instance.evaluate(solution) - current_cost
        return current_solution, improvement

    def __call__(self, solution, problem_instance, intensity=20):
        """应用相邻交换局部搜索"""
        if not self.is_applicable(problem_instance):
            return solution, 0

        # 使用高效实现
        return self.efficient_adjacent_swap(solution, problem_instance, intensity)


# 注册算子
register_operator("adjacent_swap", AdjacentSwapOperator())